import os
import pandas as pd
import numpy as np
import datetime
from matplotlib import pyplot as plt
import matplotlib.ticker as ticker
from pylab import mpl
mpl.rcParams['font.sans-serif'] = ['Microsoft YaHei']
mpl.rcParams['axes.unicode_minus'] = False
plt.switch_backend('agg')

if __name__ == '__main__':
    # 读取所有的1m K线数据
    K_LINE_FILE_PATH = "./1m/1m/"
    k_line = {}
    k_line_dates = {}
    for f in os.listdir(K_LINE_FILE_PATH):
        symbol = f.split(".")[0]
        df = pd.read_feather(K_LINE_FILE_PATH + f)
        df['trading_day'] = df['trading_day'].astype(np.str_)
        k_line.update({symbol:df.groupby(by=['trading_day'])})
        k_line_dates.update({symbol:set(df['trading_day'].tolist())}) # K线有数据的日期})
        
    # 读取目标持仓文件
    filename = 'SampleTarget2'
    SAMPLETARGET_FILE_PATH = f"./{filename}.csv"
    df_sample_target = pd.read_csv(SAMPLETARGET_FILE_PATH)
    # 数据预处理
    df_sample_target['trading_day'] = df_sample_target['trading_day'].astype(np.str_) # 日期类型转为str
    df_sample_target = df_sample_target.fillna(0) # 空值填充0
    symbol_columns = df_sample_target.columns[2:] # 品种列名
    print("K线数据读取完毕.....")
    # 初始化全部仓位大小
    if abs(df_sample_target[df_sample_target[symbol_columns].abs() > 1][symbol_columns].sum(0).sum()) > 0:
        total_positions = 1000000
    else:
        total_positions = 1 
        
    # 开始计算
    print("开始进行每日计算.....")
    fee_rate = 0.0003
    pnl_total = []
    pnlnet_total = []
    accum_fee = []
    pnl_symbol = {}
    turnover_per_day = []
    sharp_rate = 0
    turnover_rate = 0
    avg_leverage_rate = 0

    # 获取需要计算指标的日期范围
    date_format = "%Y%m%d"
    start_date_str = df_sample_target['trading_day'][0]
    end_date_str = list(df_sample_target['trading_day'])[-1]
    start_date = datetime.datetime.strptime(start_date_str, date_format)
    end_date = datetime.datetime.strptime(end_date_str, date_format)

    # 遍历每一天
    trade_records = {} # 记录每个品种的相关信息。cash:每日剩余现金， position: 仓位（以期货币种的数量记录）， accum_fee: 每日累积手续费，
                    # turnover: 每日成交金额， pnl: 每日pnl, leverage: 每日资金占有率
                    #{'a':{'cash': **, 'position': **, 'accum_fee': [**,**,...], 'turnover':[**,**,...], 'pnl':[**,**,...], 'leverage':[**,**,...]},
                    # ag':{'cash': **, 'position': **, 'accum_fee': [**,**,...], 'turnover':[**,**,...], 'pnl':[**,**,...], 'leverage':[**,**,...]},
                    # ...}
    # 初始化trade_records
    for symbol in symbol_columns:
        trade_records.update({symbol:{'cash': total_positions, 'position': 0, 'accum_fee': [], 'turnover':[], 'pnl':[], 'leverage':[]}})
        
    # 日期遍历
    date_list = []
    for i in range((end_date-start_date).days+1):
        
        today = (start_date + datetime.timedelta(days=i)).strftime(date_format)
        date_list.append(today)
        today_data = df_sample_target[df_sample_target['trading_day']==today]
        
        # 遍历品种
        for symbol in symbol_columns:
            trade_data = trade_records[symbol]
            
            # 如果没有今天的K线数据，则pnl延续上一个, 今日成交额为0
            if today not in k_line_dates[symbol]: 
                    trade_data['turnover'].append(0) # 今日成交额为0
                    trade_data['accum_fee'].append(f[-1] if (f:=trade_data['accum_fee']) else 0) # 今日累积手续费等于上一个
                    trade_data['pnl'].append(p[-1] if (p:=trade_data['pnl']) else total_positions) # 今日pnl 等于上一个
                    trade_data['leverage'].append((trade_data['pnl'][-1] - trade_data['cash'] - trade_data['accum_fee'][-1])/(trade_data['pnl'][-1] - trade_data['accum_fee'][-1])) # leverage
                    trade_records.update({symbol:trade_data})
            else:
                kline_today = k_line[symbol].get_group(today)
                #   如果没有今天的目标持仓数据，则代表仓位不变化，则今日成交额为0。
                if today_data.empty: 
                    trade_data['turnover'].append(0) # 今日成交额为0
                    trade_data['accum_fee'].append(f[-1] if (f:=trade_data['accum_fee']) else 0) # 今日累积手续费等于上一个
                    today_close_price = kline_today['clz'].tolist()[-1] # 今日收盘价
                    pnl = trade_data['accum_fee'][-1] + trade_data['cash'] + abs(trade_data['position']) * today_close_price
                    trade_data['pnl'].append(pnl)
                    trade_data['leverage'].append((trade_data['pnl'][-1] - trade_data['cash'] - trade_data['accum_fee'][-1])/(trade_data['pnl'][-1] - trade_data['accum_fee'][-1])) # leverage
                    trade_records.update({symbol:trade_data})
                else:
                    ts = today_data['timestamp'].values[0] # 买卖时间点             
                    # 改变仓位
                    symbol_target_position = today_data[symbol].values[0] # 当前品种目标仓位金额
                    price = kline_today[kline_today['timestamp']==ts]['clz'] # 成交价
                    
                    # 由于K线数据的丢失，可能在该时间点查不到clz，此时跳过本次成交，以最后一根K线计算pnl
                    if price.empty:
                        today_close_price = kline_today['clz'].tolist()[-1] # 今日收盘价
                        pnl = trade_data['accum_fee'][-1] + trade_data['cash'] + abs(trade_data['position']) * today_close_price
                        trade_data['pnl'].append(pnl)
                        trade_data['turnover'].append(0)
                        trade_data['accum_fee'].append(f[-1] if (f:=trade_data['accum_fee']) else 0) # 今日累积手续费等于上一个
                        trade_data['leverage'].append((trade_data['pnl'][-1] - trade_data['cash'] - trade_data['accum_fee'][-1])/(trade_data['pnl'][-1] - trade_data['accum_fee'][-1])) # leverage
                        trade_records.update({symbol:trade_data})
                    
                    # 成交价存在，则以该价格进行仓位调整
                    else:
                        price = price.values[0]
                        last_accum_fee = f[-1] if (f:=trade_data['accum_fee']) else 0 # 上一天的累积手续费
                        position = symbol_target_position / price # 目标仓位（以期货品种数量记）
                        last_postion = trade_data['position'] # 旧的持仓量（以期货品种数量记）
                        last_cash = trade_data['cash']
                        volume = abs(position-last_postion)# 本次需要的成交量
                        turnover = volume * price # 本次的成交额
                        fee = turnover * fee_rate # 手续费
                        if last_postion <= 0 and  position <= 0: # 之前是开空，现在仍然开空
                            if position - last_postion <= 0: # 需要开仓
                                cash = last_cash - turnover - fee # 剩余现金
                            else: # 需要平仓
                                cash = last_cash + turnover - fee # 剩余现金
                        elif last_postion <= 0 and  position >= 0: # 之前是开空，现在开多。 需要平掉旧的，开新的
                            cash = last_cash + abs(last_postion) * price - position * price - fee # 剩余现金
                        elif last_postion >= 0 and  position <= 0: # 之前是开多，现在是开空，需要平掉旧的，开心的
                            cash = last_cash + last_postion * price - abs(position) * price - fee # 剩余现金
                        elif last_postion >= 0 and  position >= 0: # 之前是凯多，现在仍然凯多
                            if position-last_postion <= 0: # 需要平仓
                                cash = last_cash + turnover - fee # 剩余现金
                            else: # 需要开仓
                                cash = last_cash - turnover - fee # 剩余现金
                                
                        accum_fee = last_accum_fee + fee # 累积手续费
                        # 收盘统计
                        today_close_price = kline_today['clz'].tolist()[-1] # 今日收盘价
                        pnl = accum_fee + cash + abs(position) * today_close_price # pnl = 现金 + 市值 (费前)
                        trade_data['cash'] = cash
                        trade_data['position'] = position
                        trade_data['accum_fee'].append(accum_fee)
                        trade_data['pnl'].append(pnl)
                        trade_data['turnover'].append(turnover) 
                        trade_data['leverage'].append((trade_data['pnl'][-1] - trade_data['cash'] - trade_data['accum_fee'][-1])/(trade_data['pnl'][-1] - trade_data['accum_fee'][-1])) # leverage
                        trade_records.update({symbol:trade_data})
                        
    # 指标统计
    print("开始进行数据指标汇总统计.....")
    df_sta = pd.DataFrame(columns=['date', 'pnl', 'pnlnet', 'total_accum_fee','turnover'])
    df_symbol_pnl = pd.DataFrame()
    df_symbol_pnlnet = pd.DataFrame()
    df_symbol_accum_fee = pd.DataFrame()
    df_symbol_trunover = pd.DataFrame()
    df_symbol_leverage = pd.DataFrame()
    for symbol, tr in trade_records.items():
        df_sta[symbol+'_pnl'] = tr['pnl'] 
        df_symbol_pnl[symbol+'_pnl'] = tr['pnl'] 
        df_symbol_trunover[symbol+'_turnover'] = tr['turnover']
        df_symbol_accum_fee[symbol+'_accumfee'] = tr['accum_fee']
        df_symbol_pnlnet[symbol+'_pnlnet'] = df_symbol_pnl[symbol+'_pnl'] - df_symbol_accum_fee[symbol+'_accumfee']
        df_symbol_leverage[symbol+'_leverage'] = tr['leverage']
        
    df_sta['total_accum_fee'] = df_symbol_accum_fee.sum(1)
    df_sta['turnover'] = df_symbol_trunover.sum(1)
    df_sta['pnl'] = df_symbol_pnl.sum(1)
    df_sta['pnlnet'] = df_sta['pnl'] - df_sta['total_accum_fee']
    df_sta['date'] = date_list

    # sharp
    # turnover
    # avgleverage
    return_rate = (df_symbol_pnlnet.iloc[-1] - total_positions) / total_positions
    risk_free_rate = 0.03
    sharp = round((np.mean(return_rate) - risk_free_rate) / np.std(return_rate), 3)
    turnover_rate = round(np.mean(df_sta['turnover']/(total_positions*len(trade_records))), 3)
    avgleverage_rate = round(np.mean(np.mean(df_symbol_leverage)), 3)
    df_symbol_pnl.insert(0, 'date', date_list)
    df_sta.to_csv(f"./result/{filename}/sta.csv", index=False)
    df_symbol_pnl.to_csv(f"./result/{filename}/pnl_each_symbol.csv", index=False)
    
    
    # 绘图
    fig = plt.figure(figsize=(20,5))
    ax = fig.add_subplot(111)
    tick_spacing = 360
    ax.xaxis.set_major_locator(ticker.MultipleLocator(tick_spacing))
    ax.plot(date_list, df_sta[['pnl','pnlnet']], label=['pnl','pnlnet'])
    ax.set_title(f"sharp_ratio:{sharp}, turnover:{turnover_rate}, avgleveage:{avgleverage_rate}")
    ax.set_xlabel("date")
    ax.legend(loc=0)
    plt.savefig(f"./result/{filename}/figure/pnl_pnlnet.png")
    plt.cla()

    # 总费用
    ax.xaxis.set_major_locator(ticker.MultipleLocator(tick_spacing))
    ax.plot(date_list, df_sta[['total_accum_fee']], label='总手续费')
    ax.set_xlabel("date")
    plt.legend(loc=0)
    plt.savefig(f"./result/{filename}/figure/total_accum_fee.png")
    plt.cla()

    # 日还手金额
    ax.xaxis.set_major_locator(ticker.MultipleLocator(tick_spacing))
    ax.plot(date_list, df_sta[['turnover']], label='日换手金额')
    ax.set_xlabel("date")
    plt.legend(loc=0)
    plt.savefig(f"./result/{filename}/figure/turnover.png")
    plt.cla()
    plt.close()